from rtree import Rtree
import geometry_algorithm
import SpatialTransform
import geometry_algorithm
import numpy as np
import ogr_load
import RoadLink
import sys

class QueryIndex(object):
    def __init__(self):
        self.is_virtual_node = False
        self.node_id = 0
        self.attach_link_id = 0
        self.attach_link = RoadLink.RoadLink()
        self.adj_nodes = []
        self.position = np.array([0.0, 0.0])

class GraphNode(object):
    def __init__(self, node):
        self.node = node
        self.edges = []

class GraphLink(object):
    def __init__(self, link):
        self.link = link

class Graph(object):
    def __init__(self, rdlinks):
        self.graph_nodes = {}
        self.used_node_list = []
        self.collected_node_dict = {}
        self.from_index = 0
        self.to_index = 0
        self.cluster_seed_index = Rtree()
        self.init_graph(rdlinks)
        self.max_node_id = self.get_max_node_id()

    def init_graph(self, rdlinks):
        for rdlink in rdlinks:
            snode = RoadLink.RoadNode()
            snode.ID = rdlink.SNID
            snode.Position = rdlink.vertexs[0]
            self.addNode(snode)

            enode = RoadLink.RoadNode()
            enode.ID = rdlink.ENID
            enode.Position = rdlink.vertexs[len(rdlink.vertexs) - 1]
            self.addNode(enode)

            link = RoadLink.RoadLink()
            link.SNID = rdlink.SNID
            link.ENID = rdlink.ENID
            link.ID = rdlink.ID
            link.vertexs = rdlink.vertexs
            link.length = link.get_length()
            self.addlink(link)

    def get_graphlinks(self):
        graphlinks = []
        for id, gnode in self.graph_nodes.items():
            for edge in gnode.edges:
               graphlinks.append(edge.link)
        return list(set(graphlinks))

    def addlink(self, link):
        self.cluster_seed_index.insert(link.ID, link.get_bound(), link)
        self.graph_nodes[link.SNID].edges.append(GraphLink(link))
        self.graph_nodes[link.ENID].edges.append(GraphLink(link))
        return True

    def addNode(self, node):
        if not node.ID in self.graph_nodes:
            self.graph_nodes[node.ID] = GraphNode(node)
        return True

    def shorest_distance(self, from_index, to_index):
        self.used_node_list = []
        self.collected_node_dict = {}
        self.collected_node_dict[from_index.node_id] = [0, -1]
        self.from_index = from_index
        self.to_index = to_index
        self.direct_distance = 0.0
        return self.get_shortest_distance()

    def shorest_path(self, from_index, to_index):
        self.used_node_list = set()
        self.collected_node_dict = {}
        self.collected_node_dict[from_index.node_id] = [0, -1]
        self.from_index = from_index
        self.to_index = to_index
        self.direct_distance = 0.0
        if self.from_index.is_virtual_node:
                self.used_node_list.append(self.from_index.node_id)
                for adj_node in self.from_index.adj_nodes:
                    distance = geometry_algorithm.calEuclideanDistance(self.graph_nodes[adj_node].node.Position, self.from_index.position)
                    self.collected_node_dict[adj_node] = [distance, self.from_index.node_id]        
        self._foreach_dijstra()
        return self._format_path()

    def get_shortest_distance(self):
        shorest_dis = 0.0
        self.direct_distance = geometry_algorithm.calEuclideanDistance(self.from_index.position, self.to_index.position)
        if self.from_index.attach_link_id == self.to_index.attach_link_id:
            shorest_dis = self.direct_distance
            return shorest_dis, list(), True
        else:
            if self.from_index.is_virtual_node:
                self.used_node_list.append(self.from_index.node_id)
                for adj_node in self.from_index.adj_nodes:
                    distance = geometry_algorithm.calEuclideanDistance(self.graph_nodes[adj_node].node.Position, self.from_index.position)
                    self.collected_node_dict[adj_node] = [distance, self.from_index.node_id]        
            self._foreach_dijstra()
            shorest_dis  = self._get_dis()
            path_and_isfind = self._format_path()
        return shorest_dis, path_and_isfind[0], path_and_isfind[1]

    def _foreach_dijstra(self):
        while True:
            if len(self.collected_node_dict) == len(self.used_node_list):
                return

            if self.to_index.node_id in self.collected_node_dict:
                return

            seed = self._find_min_f()
            self.used_node_list.append(seed)

            #great than direct twice distance, break
            if self.collected_node_dict[seed][0] > 5 * self.direct_distance:
                return

            if seed in self.to_index.adj_nodes:
                weight = self.collected_node_dict[seed][0] + geometry_algorithm.calEuclideanDistance(self.graph_nodes[seed].node.Position, self.to_index.position)
                self.collected_node_dict[self.to_index.node_id] = [weight, seed]
            else:
                adj_nodes_with_distances  = self._get_adj_nodes(seed)
                for adj_nodes_with_distance in adj_nodes_with_distances:
                    adj_node = adj_nodes_with_distance[0]
                    if adj_node not in self.used_node_list:
                        weight = self.collected_node_dict[seed][0] + adj_nodes_with_distance[1]
                        if adj_node in self.collected_node_dict:
                            if self.collected_node_dict[adj_node][0] > weight:
                                self.collected_node_dict[adj_node][0] = weight
                                self.collected_node_dict[adj_node][1] = seed
                        else:
                            self.collected_node_dict[adj_node] = [weight, seed]

    def _get_adj_nodes(self, seed_id):
        adj_nodes_with_distance = []
        for edge in self.graph_nodes[seed_id].edges:
            adj_node = edge.link.ENID if edge.link.SNID == seed_id else edge.link.SNID
            adj_nodes_with_distance.append((adj_node, edge.link.length))
        return adj_nodes_with_distance

    def _find_min_f(self):  
        min_distance = sys.float_info.max
        min_index = 0
        for k, v in self.collected_node_dict.items():
           if k in self.used_node_list:
               continue
           astar_dis = v[0] + geometry_algorithm.calEuclideanDistance(self.graph_nodes[k].node.Position, self.to_index.position)
           if astar_dis < min_distance:
               min_distance = astar_dis
               min_index = k
        return min_index

    def _get_dis(self):
        temp_node = self.to_index.node_id
        if temp_node not in self.collected_node_dict:
            return -1
        else:
            return self.collected_node_dict[temp_node][0]

    def _format_path(self):
        isfind = True
        node_list = []
        temp_node = self.to_index.node_id
        if temp_node not in self.collected_node_dict:
            isfind = False
            return node_list, isfind
        node_list.append((temp_node, self.collected_node_dict[temp_node][0]))
        while self.collected_node_dict[temp_node][1] != -1:
            temp_node = self.collected_node_dict[temp_node][1]
            node_list.append((temp_node, self.collected_node_dict[temp_node][0]))
        node_list.reverse()
        return node_list, isfind

    def get_max_node_id(self):
        maxID = -1
        for id, gnode in self.graph_nodes.items():
            if id > maxID:
                maxID = id
        return maxID

    def request_node_id(self):
        self.max_node_id = self.max_node_id + 1
        return self.max_node_id

    def get_links_by_envelope(self, coord, tolerance):
        min_x = coord[0] - tolerance
        min_y = coord[1] - tolerance
        max_x = coord[0] + tolerance
        max_y = coord[1] + tolerance

        graphlinks = [o.object for o in list(self.cluster_seed_index.intersection((min_x, min_y, max_x, max_y), True))]
        return graphlinks

    def get_match_heading(self, from_index, to_index):
        dir1 = self.to_index.position - self.from_index.position
        vertexs = self.from_index.attach_link.vertexs
        dir2 = vertexs[-1] - vertexs[0]
        dot_dir = np.dot(dir1, dir2)
        flag = True if dot_dir > 0 else False
        return flag
             
                
def search_nearest_node(lon, lat, graph, tolerance = 1e-4):
    min_distance = sys.float_info.max
    neareast_index = QueryIndex()
    nearest_link_id = 0
    nearest_path = np.array([])
    adjnodes = []

    for id, gnode in graph.graph_nodes.items():
        for edge in gnode.edges:
             ptPrj, index, distance, flag = geometry_algorithm.distance_to_linesegments(np.array([lon, lat]), edge.link.vertexs)
             if distance < min_distance:
                  min_distance = distance
                  nearest_link_id = edge.link.ID
                  adjnodes = []
                  adjnodes.append(edge.link.SNID)
                  adjnodes.append(edge.link.ENID)

    if nearest_link_id != 0 :
        neareast_index.is_virtual_node = True
        neareast_index.node_id = graph.request_node_id()
        neareast_index.attach_link_id = nearest_link_id
        neareast_index.adj_nodes = adjnodes
        neareast_index.position = np.array([lon, lat])
    return neareast_index



if __name__ == "__main__":

    map_links = ogr_load.LoadSQLiteFile("part1.sqlite")
    rdlinks = [link for k, link in map_links.items()]

    graph = Graph(rdlinks)
    
    src_index = QueryIndex()
    while src_index.node_id == 0:
        srclon = 116.43262
        srclat = 39.69295
        src_index = search_nearest_node(srclon, srclat, graph)
        if src_index.node_id == 0:
            print ("this is a invalid location, please input again")
    
    des_index = QueryIndex()
    while des_index.node_id == 0 :
        #deslon, deslat = map(float, raw_input("please input destination(116.45554 39.70784) :").split())
        deslon = 116.45554
        deslat = 39.70784
        des_index = search_nearest_node(deslon, deslat, graph)
        if des_index.node_id == 0:
            print ("this is a invalid location, please input again")

    dis, path, isfind = graph.shorest_distance(src_index, des_index)

    if not isfind:
        print ("can't find shorest path")
    else:
        print (dis)
